use linfa::{
    prelude::{SingleTargetRegression},
    traits::{Fit, Predict},
    Dataset,
};
use ndarray::{array, linalg::Dot, ArcArray2, Array, Array1, Array2, Axis};

use ndarray_rand::{rand_distr::Uniform, RandomExt};


fn main() {
    // draw_gradient();
    //test1();
    //test2();
    test_gradident();
    println!("wawa");
}
// fn draw_gradient() {
//     let plot_x = Array::random(141, Uniform::new(-1.0, 6.0));

//     // 对x进行排序
//     let mut x_vec = plot_x.to_vec();

//     x_vec.sort_by(|x, y| x.partial_cmp(y).unwrap());
//     let plot_y = x_vec
//         .iter()
//         .map(|x| (*x - 2.5f64).powi(2) - 1.0)
//         .collect::<Vec<_>>();
//     println!("{:?}", x_vec);

//     // 画图
//     // let root =
//     //     SVGBackend::new("examples/gradient_rescent/gradient.svg", (640, 480)).into_drawing_area();
//     // root.fill(&WHITE).unwrap();

//     // let mut chart = ChartBuilder::on(&root)
//     //     .x_label_area_size(40)
//     //     .y_label_area_size(40)
//     //     .build_cartesian_2d(-1.0..6.0, -2.0..10.0)
//     //     .unwrap();
//     // chart
//     //     .configure_mesh()
//     //     .disable_x_mesh()
//     //     .disable_y_mesh()
//     //     .draw()
//     //     .unwrap();
//     // chart
//     //     .draw_series(LineSeries::new(
//     //         x_vec.iter().zip(plot_y.iter()).map(|(x, y)| {
//     //             println!("{:?}", (x, y));
//     //             (*x, *y)
//     //         }),
//     //         RED.filled(),
//     //     ))
//     //     .unwrap();
// }

// fn test1() {
//     let X: Array2<f64> = 2.0 * Array::random((100, 1), Uniform::new(0.0, 1.0));
//     println!("{}", X);
//     let Y = &X * 3.0 + 4.0 + Array::random((100, 1), Uniform::new(0.0, 1.0));
//     println!("{}", Y);

//     let root =
//         SVGBackend::new("./examples/gradient_rescent/test1.svg", (640, 480)).into_drawing_area();
//     root.fill(&WHITE).unwrap();

//     let mut chart = ChartBuilder::on(&root)
//         .x_label_area_size(40)
//         .y_label_area_size(40)
//         .build_cartesian_2d(0.0..2.0, 4.0..10.0)
//         .unwrap();

//     chart
//         .configure_mesh()
//         .disable_x_mesh()
//         .disable_y_mesh()
//         .draw()
//         .unwrap();

//     chart
//         .draw_series(
//             X.index_axis(Axis(1), 0)
//                 .to_vec()
//                 .iter()
//                 .zip(Y.index_axis(Axis(1), 0).iter())
//                 .map(|(x, y)| Circle::new((*x, *y), 5, RED.filled())),
//         )
//         .unwrap();
// }

fn test2() {
    let _a = array![[1.0, 2.0], [3.0, 4.0]];
    let b = array![3.0, 4.0];
    let e = array![2.0, 3.0];
    println!("{}", b.dot(&e));
    //println!("{}", a.dot(&b.t()));
    //println!("{}", a.dot(&b));
    // println!("{:?}", b.shape());
    // println!("bbbb: {}", b.t().dot(&e));
    // let hh = b.to_shared();
    // let hh1 = hh.reshape((2, 1));
    // println!("{}", hh1.dot(&e.to_shared().reshape((1, 2))));

    let c = array![[1, 2, 3]];
    let d = c.t();
    println!("{}", d.len_of(Axis(1)));
}

fn test_gradident() {
    let X: Array2<f64> = 2.0 * Array::random((100, 10), Uniform::new(0.0, 1.0));

    let mut x_new: Array2<f64> = Array2::zeros((X.len_of(Axis(0)), 0));
    let one = Array1::ones(X.len_of(Axis(0)));
    x_new.push_column(one.view()).unwrap();
    for c in X.columns() {
        x_new.push_column(c).unwrap();
    }

    let theta: Array2<f64> = Array2::ones((11, 1));
    let mut theta = theta * 3.0;
    *theta.get_mut((0, 0)).unwrap() = 4.0;

    let Y = x_new.dot(&theta) + Array::random((100, 1), Uniform::new(0.0, 1.0));
    let initial_theta: Array2<f64> = Array2::zeros((x_new.len_of(Axis(1)), 1));
    let eta = 0.01;
    let ret = gradient_descent(
        x_new.clone(),
        Y.clone(),
        initial_theta.clone(),
        eta,
        100000,
        0.0001,
    );
    println!("{}", ret.unwrap());
    // 测试随机梯度
    let res = sgd(Dataset::new(x_new.clone(), Y.clone()), &initial_theta, 2);
    println!("res: {}", res);

    // 使用标准库
    sgd_linfa(Dataset::new(X, Y));
}

//损失函数
// X_b: 是样本的矩阵
// theta: 是系数向量
// y: 是结果向量
// 以上的向量都是竖向量
fn J(theta: &Array2<f64>, X_b: &Array2<f64>, y: &Array2<f64>) -> f64 {
    let poor = y - X_b.dot(theta);
    let poor2 = &poor * &poor;
    poor2.sum() / X_b.len_of(Axis(0)) as f64
}

// theta和y都是列向量
fn dJ(theta: &Array2<f64>, X_b: &Array2<f64>, y: &Array2<f64>) -> Array2<f64> {
    // let mut res = Array2::<f64>::zeros((theta.len_of(Axis(0)), 1));
    // *res.get_mut((0, 0)).unwrap() = (X_b.dot(theta) - y).sum();
    // for i in 1..theta.len_of(Axis(0)) {
    //     let poor = X_b.dot(theta) - y;
    //     let x = X_b.column(i).to_shared();
    //     let x = x.reshape((x.len(), 1));
    //     res.row_mut(i)[0] = *poor.t().dot(&x).get((0, 0)).unwrap();
    // }
    // println!("res: {}", res);

    // 向量运算方法
    let res = X_b.t().dot(&(X_b.dot(theta) - y));
    res * 2.0 / X_b.len_of(Axis(0)) as f64
}

// 随机梯度
// X_b_i是传进来的一行样本
fn dJ_sgd(theta: ArcArray2<f64>, X_b_i: ArcArray2<f64>, y: f64) -> Array2<f64> {
    X_b_i.dot(&(X_b_i.t().dot(&theta) - y)) * 2.0
}

// 随机梯度
// n_iters: 表示对所有的样本进行检查几轮
fn sgd(dataset: Dataset<f64, f64>, initial_theta: &Array2<f64>, n_iters: usize) -> Array2<f64> {
    let t0 = 5.0;
    let t1 = 50.0;
    let f = |t: f64| t0 / (t + t1);

    let mut theta = initial_theta.clone();
    for cur_iter in 0..n_iters {
        // 首先对所有的样本进行乱序排列
        let dataset_tmp = dataset.shuffle(&mut rand::thread_rng());
        for (idx, (r, t)) in dataset_tmp.sample_iter().enumerate() {
            let vec_r = r.to_shared().reshape((r.len(), 1));
            let vec_t = t.get(0).unwrap();
            let gradient = dJ_sgd(theta.to_shared(), vec_r, *vec_t);
            theta = theta - f((idx + dataset_tmp.targets().len() * cur_iter) as f64) * gradient;
        }
    }
    theta
}

fn gradient_descent(
    X_b: Array2<f64>,
    y: Array2<f64>,
    initial_theta: Array2<f64>,
    eta: f64,
    n_iters: usize,
    epsilon: f64,
) -> Option<Array2<f64>> {
    let mut theta = initial_theta;
    let mut i_iter = 0;
    while i_iter < n_iters {
        let gradient = dJ(&theta, &X_b, &y);
        let last_theta = theta.clone();
        theta = theta - eta * gradient;
        if (J(&theta, &X_b, &y) - J(&last_theta, &X_b, &y)).abs() < epsilon {
            return Some(theta);
        }
        i_iter += 1;
    }
    None
}

// 使用标准库调用梯度下降
fn sgd_linfa(dataset: Dataset<f64, f64>) {
    let targets_count = dataset.targets().len_of(Axis(0));
    let dataset_target_ix1 = dataset.targets().to_shared().reshape(targets_count);
    let dataset = dataset.with_targets(dataset_target_ix1);
    let model = linfa_linear::LinearRegression::default();
    let predict = model.fit(&dataset).unwrap();
    println!("{}, {}", predict.params(), predict.intercept());
    let pred = predict.predict(&dataset);
    let ret = pred.r2(&dataset).unwrap();
    println!("score: {}", ret);
    println!("std: {}", pred.std(1.0));
}
